#include <iostream>
#include <cuda_runtime.h>
#include "../common/common.h"
#define DEVICENUM 0

__host__ void summaryCPU(float *array_a, float *array_b, float *array_res, int size)
{
    for (int i = 0; i < size; i++)
        array_res[i] = array_a[i] + array_b[i];
}

__global__ void summaryGPU(float *array_a, float *array_b, float *array_res)
{
    int index = threadIdx.x + blockDim.x * blockIdx.x;
    array_res[index] = array_a[index] + array_b[index];
}

int main(int argc, char **argv)
{
    cudaSetDevice(DEVICENUM);
    int nElem = 1 << 15;
    size_t nBytes = sizeof(float) * nElem;

    double time_cpu, time_gpu;

    // 申请统一内存寻址方式的托管内存
    float *host_array_a = nullptr;
    float *host_array_b = nullptr;
    float *host_array_res = nullptr;
    float *device_array_res = nullptr;
    // 该托管内存虽然在表面上看设备和主机都能访问,但是内部其实已经自动进行了数据传输
    // 所以就不需要直径cudaMemcpy显示的进行cpu<-->数据的传输了
    CHECK(cudaMallocManaged((float **)&host_array_a, nBytes));
    CHECK(cudaMallocManaged((float **)&host_array_b, nBytes));
    CHECK(cudaMallocManaged((float **)&host_array_res, nBytes));
    CHECK(cudaMallocManaged((float **)&device_array_res, nBytes));

    // 对数组a,b进行初始化
    initialDataRandom(host_array_a, nElem);
    initialDataRandom(host_array_b, nElem);
    // 对cpu、gpu上的结果数组进行初始化
    memset(host_array_res, 0, nBytes);
    memset(device_array_res, 0, nBytes);

    // 计算cpu中两数组相加结果并计算时间
    time_cpu = get_time();
    summaryCPU(host_array_a, host_array_b, host_array_res, nElem);
    std::cout << "CPU Sum Array time:" << get_time() - time_cpu << "ms" << std::endl;

    // 调用kernel函数执行gpu数组加法运算
    dim3 block_size = block_dims(nElem);
    dim3 grid_size = grid_dims(nElem);
    // cuda warm up
    // for (size_t i = 0; i < 10; i++)
    //     summaryGPU<<<grid_size, block_size>>>(device_array_a, device_array_b, device_array_res);

    time_gpu = get_time();
    // host_array_a & host_array_b直接用即可
    summaryGPU<<<grid_size, block_size>>>(host_array_a, host_array_b, device_array_res);
    std::cout << "GPU Sum Array time:" << get_time() - time_gpu << "ms" << std::endl;
    CHECK(cudaDeviceSynchronize());

    // 将gpu运算结果与cpu结果进行比较
    checkResult(host_array_res, device_array_res, nElem);

    // 释放托管内存
    cudaFree(host_array_a);
    cudaFree(host_array_b);
    cudaFree(host_array_res);
    cudaFree(device_array_res);

    return 0;
}